
import os
import ssl
import time
# from concurrent.futures import ThreadPoolExecutor
#
# import requests
from locust import task, TaskSet
from locust.user.wait_time import between
from locust_plugins.users import MqttUser
# from common.utils import enconde_bs64
# from common import logger
from testcases.api_stress.mqtt_demo.baseScript import login, mqtt_conf
# tls_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
# tls_context.load_verify_locations(os.environ["LOCUST_MQTT_CAFILE"])

# aa = enconde_bs64('12345')
# print(aa)
cookie = login('cc1234')
mqtt_con = mqtt_conf(cookie)
appGuestSendTopic = mqtt_con['appGuestSendTopic']
appGuestReceiveTopic = mqtt_con['appGuestReceiveTopic']
# def mqtt_conf(cookie):
#     """
#     去登录接口获取cookie信息,调用该接口
#     :param username: 用户名
#     :param passwd: 密码
#     :return: 返回mqtt配置信息
#     """
#     mqtt_url = "https://www.retone-hearing.fit" + '/platformRetoneFit/api/mqtt/needAppUserAuth/responseEncode/requestDecode/getMqttConnectProperties'
#     info = {}
#     # cookie = login(username, passwd)
#     res = requests.post(mqtt_url, info, cookie)
#     if res.status_code != 200 or res.json()['data'] is None:
#         logger.error('===== 获取mqtt配置信息接口失败 =========')
#         logger.error(res.json())
#     data_str = decode64_str(res.json()['data'])
#     data = json.loads(data_str)
#     dic_conf = {
#         "appGuestPassword": "123456@appGuest.retone",
#         "port": data['port'],
#         "appGuestSendTopic": data['appGuestSendTopic'],
#         "appGuestReceiveTopic": data['appGuestReceiveTopic'],
#         "appGuestReceiveClientId": data['appGuestReceiveClientId'],
#         "url": data['url'],
#         "appGuestName": "appGuest"
#     }
#     # logger.error(dic_conf)
#     return dic_conf
#
#
# l = []
# with open("../locustfiles/cookies.txt") as f:
#     for account in f:
#         l.append(account.strip())
# with ThreadPoolExecutor(50) as t:
#     # l 列表控制登录用户数量
#     for username in l[:50]:
#         future = t.submit(mqtt_conf, username.strip())

class MyUser(MqttUser):
    # get_conf = mqtt_conf(cookie)
    # logger.info('Cookie = ', cookie, 'get_conf = ', get_conf)
    broker = mqtt_con['url']
    port = 1883
    client_id = mqtt_con['appGuestReceiveClientId']
    password = mqtt_con['appGuestPassword']
    username = mqtt_con['appGuestName']


    # tls_context = tls_context
    # We could uncomment below to use the WebSockets transport
    # transport = "websockets"
    # ws_path = "/mqtt/custom/path"

    # We'll probably want to throttle our publishing a bit: let's limit it to
    # 10-100 messages per second.
    wait_time = between(0.01, 0.1)

    @task
    class MyTasks(TaskSet):
        # Sleep for a while to allow the client time to connect.
        # This is probably not the most "correct" way to do this: a better method
        # might be to add a gevent.event.Event to the MqttClient's on_connect
        # callback and wait for that (with a timeout) here.
        # However, this works well enough for the sake of an example.
        def on_start(self):
            time.sleep(5)

        @task
        def say_hello(self):
            # self.client.publish("hello/locust", b"hello world")
            i = 1
            while 1:
                time.sleep(15)
                msg = {
                    "token": cookie,
                    "messageType": "APP_SEND_SELF_ADJUSTMENT_PARAM",
                    "msgBody": {
                        "earType": "LEFT",
                        "xAxis": 0 if i % 2 == 0 else 1,
                        "yAxis": 1,
                        "programNumber": 0
                    },
                    "sendTime": "2021-04-07 08:06:10"
                }
                # logger.info(client.is_connected())
                import json
                msg = json.dumps(msg)
                result = self.client.publish(appGuestSendTopic, msg)
                # result: [0, 1]
                # status = result[0]
                # if status == 0:
                print(f"Send '{msg}' to topic '{appGuestSendTopic}'")
                # else:
                print(f"Failed to send message to topic '{appGuestSendTopic}'")
                i += 1

        @task
        def subscribe(self):
            def on_message(self, userdata, msg):
                print(f"Received `{msg.payload.decode()}` from `{msg.topic}` topic")
            self.client.subscribe(appGuestReceiveTopic)
            # logger.info(appGuestReceiveTopic)
            self.client.on_message = on_message



if __name__ == '__main__':
    os.system('locust -f locust-demo-mqtt.py --web-host=127.0.0.1')
